/*** Created Time: 2022年5月20日 18:31:17 ***/
/*** Author: Eric                        ***/

#include <onix/memory.h>
#include <onix/types.h>
#include <onix/debug.h>
#include <onix/assert.h>
#include <onix/stdlib.h>
#include <onix/string.h>
#include <onix/bitmap.h>
#include <onix/multiboot2.h>
#include <onix/task.h>

// 辅助函数声明
static u32 get_page(void);
static void put_page(u32 addr);
static _inline void enable_page(void);
static void entry_init(page_entry_t *entry, u32 index);
static page_entry_t *get_pde(void);
static page_entry_t *get_pte(u32 vaddr, bool create);

#define     LOGK(fmt, args...)  DEBUGK(fmt, ##args)

#define     ZONE_VALID          (1)     // ards 可用内存区域
#define     ZONE_RESERVED       (2)     // ards 不可用区域

#define     IDX(addr)           ((u32)addr >> 12)   // 获取addr的页索引
                                                    // 右移12位的本质是屏蔽一个页大小0x1000的后面的0
#define     DIDX(addr)          (((u32)addr >> 22) & 0x3ff)     // 获取addr的页目录索引，最高的10位
#define     TIDX(addr)          (((u32)addr >> 12) & 0x3ff)     // 获取addr的页表索引，中间的10位
#define     PAGE(idx)           ((u32)idx << 12)                // 获取页索引idx对应的页开始的位置
#define     ASSERT_PAGE(addr)   assert((addr & 0xfff) == 0)

#define     PDE_MASK            (0xFFC00000)

// 内核页表索引，两张页表总共映射8M的内核内存
static u32 KERNEL_PAGE_TABLE[] = {
    0x2000,
    0x3000,
};

// 内核虚拟内存位图存到0x4000的位置，即内核位图的存储位置
#define     KERNEL_MAP_BITS     (0x4000)

// 声明内核位图
bitmap_t kernel_map;

// Address Descriptor Structure
typedef struct ards_t {
    u64 base;   // 内存基地址
    u64 size;   // 内存长度
    u32 type;   // 类型
} _packed ards_t;

static u32 memory_base = 0;    // 可用内存基地址，应该等于1M
static u32 memory_size = 0;    // 可用内存大小
static u32 total_pages = 0;    // 所有内存页数
static u32 free_pages = 0;     // 空闲内存页数

// 已用页数
#define     used_pages          (total_pages - free_pages) 

// @magic : 内核魔数
// @addr : ards_count，是ards的数量指针
void memory_init(u32 magic, u32 addr)
{
    u32 count = 0;

    // 如果是从loader.asm文件进入的内核
    if(magic == ONIX_MAGIC) {
        count = *(u32 *)addr;
        ards_t *ptr = (ards_t *)(addr + 4);

        for(size_t i=0; i<count; i++,ptr++) {
            LOGK("Memory base 0x%p  size 0x%p  type %d\n",
                 (u32)ptr->base, (u32)ptr->size, (u32)ptr->type);
            // 如果该块内存属于可用区域，并且大小大于目前可用内存大小
            if(ptr->type == ZONE_VALID && ptr->size > memory_size) {
                memory_base = (u32)ptr->base;
                memory_size = (u32)ptr->size;
            }
        }
    } else if(magic == MULTIBOOT2_MAGIC) {   // GRUB的magic：0x36d76289
        u32 size = *(unsigned int *)addr;   // 
        multi_tag_t *tag = (multi_tag_t *)(addr + 8);

        LOGK("Announced mbi size 0x%x\n", size);
        while (tag->type != MULTIBOOT_TAG_TYPE_END) {
            if (tag->type == MULTIBOOT_TAG_TYPE_MMAP) {
                break;
            }

            // 下一个 tag 对齐到了 8 字节
            tag = (multi_tag_t *)((u32)tag + ((tag->size + 7) & ~7));
        }

        // 进入到了multiboot mmap tag
        multi_tag_mmap_t *mtag = (multi_tag_mmap_t *)tag;
        multi_mmap_entry_t *entry = mtag->entries;
        while ((u32)entry < (u32)tag + tag->size)
        {
            LOGK("Memory base 0x%p size 0x%p type %d\n",
                 (u32)entry->addr, (u32)entry->len, (u32)entry->type);

            count++;

            if (entry->type == ZONE_VALID && entry->len > memory_size) {
                memory_base = (u32)entry->addr;     // 基地址
                memory_size = (u32)entry->len;      // 尺寸
            }

            entry = (multi_mmap_entry_t *)((u32)entry + mtag->entry_size);
        }
    } else {
        panic("Memory init magic unknown 0x%p\n", magic);
    }

    LOGK("ARDS count %d\n", count);                 // 打印ARDS数量
    LOGK("Memory base 0x%p\n", (u32)memory_base);   // 打印可用内存基地址
    LOGK("Memory size 0x%p\n", (u32)memory_size);   // 打印可用内存大小

    assert(memory_base == MEMORY_BASE);             // 内存开始的位置为1M
    assert((memory_size & 0xfff) == 0);             // 要求按页（4K，0x1000）对齐

    total_pages = IDX(memory_size) + IDX(MEMORY_BASE);  // 所有内存页数
    free_pages = IDX(memory_size);                      // 所有空闲页数

    LOGK("Total pages %d\n", total_pages);          // 打印所有内存页数
    LOGK("Free pages %d\n", free_pages);            // 打印所有空闲页数

    // 如果可用内存大小小于内核所需的8M内存，直接报错
    if(memory_size < KERNEL_MEMORY_SIZE) {
        panic("System memory is %dM too small, at least %dM needed\n",
                memory_size / MEMORY_BASE, KERNEL_MEMORY_SIZE / MEMORY_BASE);
    }
}


static u32 start_page = 0;      // 可分配物理内存起始地址
static u8 *memory_map;          // 物理内存数组,记录引用数量，8位可最大记录255
static u32 memory_map_pages;    // 物理内存数组占用的页数

// 初始化物理内存
void memory_map_init(void)
{
    // 初始化物理内存数组，从1M的物理地址开始
    memory_map = (u8 *)memory_base;

    // 计算物理内存数组占用的页数
    memory_map_pages = div_round_up(total_pages, PAGE_SIZE);
    LOGK("Memory map page count %d\n", memory_map_pages);

    free_pages -= memory_map_pages; // 求可用空闲内存页数

    // 清空物理内存数组
    memset((void *)memory_map, 0, memory_map_pages * PAGE_SIZE);

    // 统计前1M的内存位置以及物理内存数组已占用的页，计算出可分配物理内存的起始地址
    // 注：因为最多有4G的物理空间，一个页面4k，因此需要4G/4K=1M个页来管理物理空间
    //     即1M个页表项
    start_page = IDX(MEMORY_BASE) + memory_map_pages;
    for(size_t i=0; i<start_page; i++) {
        memory_map[i] = 1;  // 初始化该页面的被引用数量为1，可理解为有1个进程引用了该页
    }

    LOGK("Total pages %d free pages %d\n", total_pages, free_pages);    // 打印

    // 初始化内核虚拟内存位图，需要8位对齐
    u32 length = (IDX(KERNEL_MEMORY_SIZE) - IDX(memory_base)) / 8;
    bitmap_init(&kernel_map, (u8 *)KERNEL_MAP_BITS, length, IDX(memory_base));
    bitmap_scan(&kernel_map, memory_map_pages);     // 将物理内存占用的页数的位图标记为占用
}

// 分配一页物理内存
static u32 get_page(void)
{
    for(size_t i=start_page; i<total_pages; i++) {
        // 如果有一页物理内存没有被占用
        if(!memory_map[i]) {
            memory_map[i] = 1;  // 设置有一个进程占用该页

            --free_pages;       // 空闲内存页数自减
            assert(free_pages >= 0);

            u32 page = ((u32)i) << 12;      // 求得被分配的页的开始地址
            LOGK("GET page 0x%p\n", page);  // 打印被分配的内存地址

            return page;
        }
    }

    panic("Out of Memory!!!\n");    // 分配失败时运行该函数
}

// 进程主动释放一页物理内存
static void put_page(u32 addr)
{
    ASSERT_PAGE(addr);  // 检查addr是否是4K对齐

    u32 idx = IDX(addr);    // 求得addr的页索引

    // idx 大于 1M 并且小于总页面数
    assert(idx >= start_page && idx < total_pages);

    // 确保待释放的页面有被进程所使用
    assert(memory_map[idx] >= 1);

    // 物理引用减一
    memory_map[idx]--;

    // 如果引用数为0，则空闲页面加一
    if(!memory_map[idx]) {
        free_pages++;
    }

    assert(free_pages > 0 && free_pages < total_pages);
    LOGK("PUT page 0x%p\n", addr);
}

// 得到 cr2 寄存器
u32 get_cr2(void)
{
    // 直接将 mov eax, cr2，返回值在 eax 中
    __asm__ __volatile__("movl %cr2, %eax\n");
}

// 得到CR3寄存器的值
u32 get_cr3(void)
{
    // 直接 mov eax, cr3 ,返回值保存在 eax 寄存器中
    __asm__ __volatile__("movl %cr3, %eax\n");
}

// 设置CR3寄存器，参数是页目录的地址
void set_cr3(u32 pde)
{
    ASSERT_PAGE(pde);   // 检查地址是否有效4K对齐
    // 在扩展的GCC内联汇编中，寄存器需要用两个'%%'来表示
    // 约束"a"表示使用eax及相关寄存器，在执行movl指令前，将pde的值移到eax寄存器中
    __asm__ __volatile__("movl %%eax, %%cr3\n" :: "a"(pde));
}

// 将CR0寄存器最高位PG置为1，启用分页
static _inline void enable_page(void)
{
    // 0xb0000000 = 
    // 0b1000_0000_0000_0000_0000_0000_0000_0000
    __asm__ __volatile__(
        "movl %cr0, %eax\n"
        "orl $0x80000000, %eax\n"
        "movl %eax, %cr0\n"
    );
}

// 初始化页表项
static void entry_init(page_entry_t *entry, u32 index)
{
    *(u32 *)entry = 0;      // 将entry结构体初始化为0
    entry->present = 1;     // 表示在内存中
    entry->write = 1;       // 可读写
    entry->user = 1;        // 所有人可访问
    entry->index = index;   // 20位的页索引
}

// 初始化内存映射
void mapping_init(void)
{
    page_entry_t *pde = (page_entry_t *)KERNEL_PAGE_DIR;    // 将内核页目录作为起始地址(页目录表)
    memset((void *)pde, 0, PAGE_SIZE);

    idx_t index = 0;

    // sizeof(KERNEL_PAGE_TABLE)/4 = 2
    for(idx_t didx=0; didx<(sizeof(KERNEL_PAGE_TABLE)/4); didx++) {
        page_entry_t *pte = (page_entry_t *)KERNEL_PAGE_TABLE[didx];
        memset(pte, 0, PAGE_SIZE);

        page_entry_t *dentry = &pde[didx];      // 将页目录表的前两项映射到两张页表上
        entry_init(dentry, IDX((u32)pte));

        // 映射8M页表空间
        for(size_t tidx=0; tidx<1024; tidx++,index++) {
            // 第0页不映射（作为nullptr），为造成空指针访问，缺页异常，便于排错
            if(index == 0) {
                continue;
            }

            page_entry_t *tentry = &pte[tidx];  // 映射页表的页表项对应到内存空间
            entry_init(tentry, index);
            memory_map[index] = 1;              // 设置物理内存数组，该页被占用
        }
    }

    // 将最后一个页表指向页目录自己，方便修改
    page_entry_t *entry = &pde[1023];
    entry_init(entry, IDX(KERNEL_PAGE_DIR));

    // 设置CR3寄存器，将pde的首地址作为页目录表的地址
    set_cr3((u32)pde);

    // 分页有效
    enable_page();
}

// 求页目录（这里求最后一个页目录项的地址）
static page_entry_t *get_pde(void)
{
    return (page_entry_t *)(0xfffff000);
}

// 获取虚拟地址 vaddr 对应的页表
static page_entry_t *get_pte(u32 vaddr, bool create)
{
    page_entry_t *pde = get_pde();
    u32 idx = DIDX(vaddr);
    page_entry_t *entry = &pde[idx];

    // create 为 true 表示页表不存在，需要创建
    // create 为 false 表示已有页表
    assert(create || (!create && entry->present));

    // 求可修改的页表地址
    page_entry_t *table = (page_entry_t *)(PDE_MASK | (idx << 12));

    // 如果页表不存在
    if (!entry->present) {
        LOGK("Get and create page table entry for 0x%p\n", vaddr);
        u32 page = get_page();          // 初始化一个新的页面
        entry_init(entry, IDX(page));
        memset(table, 0, PAGE_SIZE);
    }

    return table;
}

// 刷新虚拟地址vaddr的快表TLB
// r 表示任何可用的通用寄存器，%0表示vaddr
static void flush_tlb(u32 vaddr)
{
    asm volatile("invlpg (%0)" ::"r"(vaddr)
                 : "memory");
}

// 从位图中扫描count个连续的页
static u32 scan_page(bitmap_t *map, u32 count)
{
    assert(count > 0);
    int32 index = bitmap_scan(map, count);

    if (index == EOF) {
        panic("Scan page fail!!!");
    }

    u32 addr = PAGE(index);
    LOGK("Scan page 0x%p count %d\n", addr, count);
    return addr;
}

// 与scan_page相对，重置相应的页
static void reset_page(bitmap_t *map, u32 addr, u32 count)
{
    ASSERT_PAGE(addr);
    assert(count > 0);
    u32 index = IDX(addr);

    for (size_t i = 0; i < count; i++) {
        assert(bitmap_test(map, index + i));
        bitmap_set(map, index + i, 0);
    }
}

// 分配count个连续的内核页
u32 alloc_kpage(u32 count)
{
    assert(count > 0);
    u32 vaddr = scan_page(&kernel_map, count);
    LOGK("ALLOC kernel pages 0x%p count %d\n", vaddr, count);
    return vaddr;
}

// 释放count个连续的内核页
void free_kpage(u32 vaddr, u32 count)
{
    ASSERT_PAGE(vaddr);
    assert(count > 0);
    reset_page(&kernel_map, vaddr, count);
    LOGK("FREE  kernel pages 0x%p count %d\n", vaddr, count);
}

// 将 vaddr 映射物理内存
void link_page(u32 vaddr)
{
    ASSERT_PAGE(vaddr);     // 确保 vaddr 是页开始的位置

    page_entry_t *pte = get_pte(vaddr, true);   // 得到对应的页表
    page_entry_t *entry = &pte[TIDX(vaddr)];    // 得到页框对应的入口地址

    task_t *task = running_task();
    bitmap_t *map = task->vmap;
    u32 index = IDX(vaddr);

    // 如果页面已存在，则直接返回
    if (entry->present) {
        assert(bitmap_test(map, index));
        return;
    }

    assert(!bitmap_test(map, index));
    bitmap_set(map, index, true);

    u32 paddr = get_page();
    entry_init(entry, IDX(paddr));
    flush_tlb(vaddr);

    LOGK("LINK from 0x%p to 0x%p\n", vaddr, paddr);
}

// 去掉 vaddr 对应的物理内存映射
void unlink_page(u32 vaddr)
{
    ASSERT_PAGE(vaddr);

    page_entry_t *pte = get_pte(vaddr, true);
    page_entry_t *entry = &pte[TIDX(vaddr)];

    task_t *task = running_task();
    bitmap_t *map = task->vmap;
    u32 index = IDX(vaddr);

    // 如果页面不在内存中
    if (!entry->present) {
        assert(!bitmap_test(map, index));
        return;
    }

    assert(entry->present && bitmap_test(map, index));

    entry->present = false;
    bitmap_set(map, index, false);

    u32 paddr = PAGE(entry->index);

    DEBUGK("UNLINK from 0x%p to 0x%p\n", vaddr, paddr);
    put_page(paddr);    // 释放页面

    flush_tlb(vaddr);   // 刷新 TLB
}

// 拷贝一页，返回拷贝后的物理地址
static u32 copy_page(void *page)
{
    u32 paddr = get_page();

    page_entry_t *entry = get_pte(0, false);
    entry_init(entry, IDX(paddr));
    memcpy((void *)0, (void *)page, PAGE_SIZE);

    entry->present = false;
    return paddr;
}

// 拷贝当前页目录
page_entry_t *copy_pde(void)
{
    // 拷贝当前进程的页表
    task_t *task = running_task();

    // 申请新的一页内存，存放当前进程的页表
    page_entry_t *pde = (page_entry_t *)alloc_kpage(1); // todo free
    memcpy(pde, (void *)task->pde, PAGE_SIZE);

    // 将最后一个页表指向页目录自己，方便修改
    page_entry_t *entry = &pde[1023];
    entry_init(entry, IDX(pde));

    page_entry_t *dentry;

    for (size_t didx = 2; didx < 1023; didx++) {
        dentry = &pde[didx];
        if (!dentry->present)
            continue;

        page_entry_t *pte = (page_entry_t *)(PDE_MASK | (didx << 12));

        for (size_t tidx = 0; tidx < 1024; tidx++) {
            entry = &pte[tidx];
            if (!entry->present)
                continue;

            // 对应物理内存引用大于 0
            assert(memory_map[entry->index] > 0);
            // 置为只读
            entry->write = false;
            // 对应物理页引用加 1
            memory_map[entry->index]++;

            assert(memory_map[entry->index] < 255);
        }

        u32 paddr = copy_page(pte);
        dentry->index = IDX(paddr);
    }

    set_cr3(task->pde);

    return pde;
}

// 释放当前页目录
void free_pde(void)
{
    task_t *task = running_task();
    assert(task->uid != KERNEL_USER);

    page_entry_t *pde = get_pde();

    for (size_t didx = 2; didx < 1023; didx++) {
        page_entry_t *dentry = &pde[didx];
        if (!dentry->present) {
            continue;
        }

        page_entry_t *pte = (page_entry_t *)(PDE_MASK | (didx << 12));

        for (size_t tidx = 0; tidx < 1024; tidx++) {
            page_entry_t *entry = &pte[tidx];
            if (!entry->present) {
                continue;
            }

            assert(memory_map[entry->index] > 0);
            put_page(PAGE(entry->index));
        }

        // 释放页表
        put_page(PAGE(dentry->index));
    }

    // 释放页目录
    free_kpage(task->pde, 1);
    LOGK("free pages %d\n", free_pages);
}

int32 sys_brk(void *addr)
{
    LOGK("task brk 0x%p\n", addr);
    u32 brk = (u32)addr;
    ASSERT_PAGE(brk);   // 传进来的 addr 应该是页开始的位置

    task_t *task = running_task();
    assert(task->uid != KERNEL_USER);   // 确保是用户态调用

    assert(KERNEL_MEMORY_SIZE < brk < USER_STACK_BOTTOM);

    u32 old_brk = task->brk;    // 原来的 brk 地址

    if (old_brk > brk)
    {
        for (; brk < old_brk; brk += PAGE_SIZE)
        {
            unlink_page(brk);   // 如果原来的 brk 地址大于现在的地址，释放内存
        }
    }
    else if (IDX(brk - old_brk) > free_pages)
    {
        // out of memory
        return -1;
    }

    task->brk = brk;
    return 0;
}

typedef struct page_error_code_t
{
    u8 present : 1;
    u8 write : 1;
    u8 user : 1;
    u8 reserved0 : 1;
    u8 fetch : 1;
    u8 protection : 1;
    u8 shadow : 1;
    u16 reserved1 : 8;
    u8 sgx : 1;
    u16 reserved2;
} _packed page_error_code_t;

void page_fault(
    u32 vector,
    u32 edi, u32 esi, u32 ebp, u32 esp,
    u32 ebx, u32 edx, u32 ecx, u32 eax,
    u32 gs, u32 fs, u32 es, u32 ds,
    u32 vector0, u32 error, u32 eip, u32 cs, u32 eflags)
{
    assert(vector == 0xe);  // 确认缺页异常
    u32 vaddr = get_cr2();  // 从 cr2 寄存器中得到引发缺页异常的地址
    LOGK("fault address 0x%p\n", vaddr);

    page_error_code_t *code = (page_error_code_t *)&error;
    task_t *task = running_task();

    // 确保 vaddr 在用户态的内存空间中
    assert(KERNEL_MEMORY_SIZE <= vaddr < USER_STACK_TOP);

    if (code->present) {
        assert(code->write);

        page_entry_t *pte = get_pte(vaddr, false);
        page_entry_t *entry = &pte[TIDX(vaddr)];

        assert(entry->present);
        assert(memory_map[entry->index] > 0);
        if (memory_map[entry->index] == 1) {
            entry->write = true;
            LOGK("WRITE page for 0x%p\n", vaddr);
        } else {
            void *page = (void *)PAGE(IDX(vaddr));
            u32 paddr = copy_page(page);
            memory_map[entry->index]--;
            entry_init(entry, IDX(paddr));
            flush_tlb(vaddr);
            LOGK("COPY page for 0x%p\n", vaddr);
        }
        return;
    }

    if (!code->present && (vaddr < task->brk || vaddr >= USER_STACK_BOTTOM))
    {
        u32 page = PAGE(IDX(vaddr));
        link_page(page);
        // BMB;
        return;
    }

    panic("page fault!!!");
}